﻿#include "FileMesh.h"
#include "Logging.h"
#include "Timer.h"
#include <fstream>
#include <sstream>

#include <vtkCellData.h>
#include <vtkCellType.h>
#include <vtkCompositeDataSet.h>
#include <vtkDataArraySelection.h>
#include <vtkIOSSReader.h>
#include <vtkIdList.h>
#include <vtkInformation.h>
#include <vtkInformationIntegerKey.h>
#include <vtkInformationKey.h>
#include <vtkMergeCells.h>
#include <vtkPartitionedDataSet.h>
#include <vtkPartitionedDataSetCollection.h>
#include <vtkPointData.h>
#include <vtkSmartPointer.h>
#include <vtkUnstructuredGrid.h>
#include <vtkmlib/DataSetConverters.h>
#include <vtkmlib/UnstructuredGridConverter.h>

#include <vtkm/io/VTKDataSetWriter.h>


RegisterObject(FileMesh);
void FileMesh::Build()
{
  ReadExodusMesh();
  Prepared();
}

vtkUnstructuredGrid* FileMesh::MergeGrids(vtkDataSet** sets,
                                          int nsets,
                                          int deleteDataSets,
                                          int useGlobalNodeIds,
                                          float pointMergeTolerance,
                                          int useGlobalCellIds)
{
  int i;

  if (nsets == 0)
  {
    return nullptr;
  }

  vtkUnstructuredGrid* newGrid = vtkUnstructuredGrid::New();
  // Any global ids should be consistent, so make sure they are passed.
  newGrid->GetPointData()->CopyGlobalIdsOn();
  newGrid->GetCellData()->CopyGlobalIdsOn();

  vtkMergeCells* mc = vtkMergeCells::New();
  mc->SetUnstructuredGrid(newGrid);

  mc->SetTotalNumberOfDataSets(nsets);

  vtkIdType totalPoints = 0;
  vtkIdType totalCells = 0;

  for (i = 0; i < nsets; i++)
  {
    totalPoints += sets[i]->GetNumberOfPoints();
    totalCells += sets[i]->GetNumberOfCells();
    // Only use global ids if they are available.
    useGlobalNodeIds = (useGlobalNodeIds && (sets[i]->GetPointData()->GetGlobalIds() != nullptr));
    useGlobalCellIds = (useGlobalCellIds && (sets[i]->GetCellData()->GetGlobalIds() != nullptr));
  }

  mc->SetTotalNumberOfPoints(totalPoints);
  mc->SetTotalNumberOfCells(totalCells);

  if (!useGlobalNodeIds)
  {
    mc->SetPointMergeTolerance(pointMergeTolerance);
  }
  mc->SetUseGlobalIds(useGlobalNodeIds);
  mc->SetUseGlobalCellIds(useGlobalCellIds);

  for (i = 0; i < nsets; i++)
  {
    mc->MergeDataSet(sets[i]);

    if (deleteDataSets)
    {
      sets[i]->Delete();
    }
  }

  mc->Finish();
  mc->Delete();

  return newGrid;
}

int FileMesh::ReadExodusMesh()
{
  // vtkIossreader read element block and side set
  auto rdr = vtkSmartPointer<vtkIOSSReader>::New();
  rdr->SetFileName(_mesh_file.c_str());
  rdr->UpdateInformation();
  rdr->GetElementBlockSelection()->EnableAllArrays(); //enable element block
  rdr->GetSideSetSelection()->EnableAllArrays();      // enable side set
  rdr->Update();

  //get entity map
  std::map<std::string, vtkIOSSReader::EntityType> entity_map;
  auto block_selection = rdr->GetElementBlockSelection();
  auto sideset_selection = rdr->GetSideSetSelection();
  auto block_number = block_selection->GetNumberOfArrays();
  auto sideset_number = sideset_selection->GetNumberOfArrays();

  for (int i_block = 0; i_block < block_number; ++i_block)
  {
    //insert element block
    entity_map.insert(
      std::make_pair(block_selection->GetArrayName(i_block), vtkIOSSReader::ELEMENTBLOCK));
  }

  for (int i_sideset = 0; i_sideset < sideset_number; ++i_sideset)
  {
    //insert side set
    entity_map.insert(
      std::make_pair(sideset_selection->GetArrayName(i_sideset), vtkIOSSReader::SIDESET));
  }
  if (0 >= entity_map.size())
  {
    return -1;
  }

  //merge element block and side set
  vtkDataSet** block_grids = new vtkDataSet*[block_number];     //for merge block
  vtkDataSet** sideset_grids = new vtkDataSet*[sideset_number]; //for merge block
  auto block_index = 0;
  auto sideset_index = 0;
  auto id_index = 0;

  auto collection = vtkPartitionedDataSetCollection::SafeDownCast(rdr->GetOutputDataObject(0));
  if (nullptr == collection)
  {
    return -1;
  }

  auto collection_number = collection->GetNumberOfPartitionedDataSets();
  for (auto i_part = 0; i_part < collection_number; ++i_part)
  {
    auto pd = collection->GetPartitionedDataSet(i_part);
    if (nullptr == pd)
    {
      continue;
    }

    auto data_set = pd->GetPartition(0);
    if (nullptr == data_set)
    {
      continue;
    }

    auto part_name = collection->GetMetaData(i_part)->Get(vtkCompositeDataSet::NAME());
    auto iter_part = entity_map.find(part_name);
    if (entity_map.end() == iter_part)
    {
      continue;
    }


    //add id field
    auto block_id_array = vtkSmartPointer<vtkIntArray>::New();
    block_id_array->SetName("cqai_cell_id");
    block_id_array->SetNumberOfComponents(1);
    auto cell_number = data_set->GetNumberOfCells();
    block_id_array->SetNumberOfTuples(cell_number);
    for (int i_cell = 0; i_cell < cell_number; ++i_cell)
    {
      block_id_array->SetTuple1(i_cell, id_index);
    }
    data_set->GetCellData()->AddArray(block_id_array);

    //add grid
    auto entity_type = iter_part->second;
    switch (entity_type)
    {
      case vtkIOSSReader::ELEMENTBLOCK:
        block_grids[block_index] = data_set;
        ++block_index;
        _id.insert(std::make_pair(part_name, id_index));
        ++id_index;
        break;
      case vtkIOSSReader::SIDESET:
        sideset_grids[sideset_index] = data_set;
        ++sideset_index;
        _id.insert(std::make_pair(part_name, id_index));
        ++id_index;
        FillElementSides(part_name, data_set);
        break;
      default:
        break;
    }
  }

  if (block_index != block_number || sideset_index != sideset_number)
  {
    return -1;
  }

  //convert vtkUnStructureGrid to vtkm dataset
  auto _block_unstructuregrid = MergeGrids(block_grids, block_number, 0, 1, 0.0, 1);
  if (nullptr == _block_unstructuregrid)
  {
    return -1;
  }

  _data_set = tovtkm::Convert(_block_unstructuregrid, tovtkm::FieldsFlag::PointsAndCells);

  auto _sideset_unstructuregrid = MergeGrids(sideset_grids, sideset_number, 0, 1, 0.0, 1);
  if (nullptr == _sideset_unstructuregrid)
  {
    return -1;
  }

  //_sideset_dataset = tovtkm::Convert(_sideset_unstructuregrid, tovtkm::FieldsFlag::PointsAndCells);

  return 0;
}


void FileMesh::FillElementSides(const std::string name, vtkDataSet* const& data_set)
{
  if (name.empty() || nullptr == data_set)
  {
    return;
  }

  //get element side
  auto cell_number = data_set->GetNumberOfCells();
  auto cell_field = data_set->GetCellData();
  auto elementside_array = cell_field->GetArray("element_side");
  auto num_comp = elementside_array->GetNumberOfComponents();
  if (2 != num_comp || nullptr == elementside_array)
  {
    return;
  }

  std::vector<FaceInfo> element_side;
  for (auto i_cell = 0; i_cell < cell_number; ++i_cell)
  {
    UInt8 edgeIndex = elementside_array->GetTuple(i_cell)[1] - 1; //edge ids -1
    Id id = elementside_array->GetTuple(i_cell)[0] - 1;           //id

    element_side.emplace_back(id, edgeIndex);
  }

  _element_sides.insert({ name, element_side });
}